{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A Case Study - Seizure Detection"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Seizure detection dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import memtorch\n",
    "from torch.utils.data import Dataset\n",
    "import torch.nn.functional as F\n",
    "import torch.nn as nn\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import sklearn\n",
    "from sklearn import preprocessing\n",
    "\n",
    "\n",
    "class SeizureDataset(Dataset):\n",
    "\n",
    "    def __init__(self, path_to_csv):\n",
    "        self.features = pd.read_csv(path_to_csv)\n",
    "        self.labels = self.features.pop('y')\n",
    "        self.features = preprocessing.scale(self.features.iloc[:, 1:], axis=0)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.labels)\n",
    "\n",
    "    def __getitem__(self, i):\n",
    "        if self.labels[i] == 1:\n",
    "            label = 1\n",
    "        else:\n",
    "            label = 0\n",
    "\n",
    "        return np.asarray(self.features[i, :]).astype(np.float), label\n",
    "\n",
    "csv_path = 'http://archive.ics.uci.edu/ml/machine-learning-databases/00388/data.csv'\n",
    "dataset = SeizureDataset(path_to_csv=csv_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Network architecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EEGNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(EEGNet, self).__init__()\n",
    "        self.fc1 = nn.Linear(178, 200)\n",
    "        self.batchnorm1 = nn.BatchNorm1d(200)\n",
    "        self.fc2 = nn.Linear(200, 200)\n",
    "        self.batchnorm2 = nn.BatchNorm1d(200)\n",
    "        self.fc3 = nn.Linear(200, 200)\n",
    "        self.batchnorm3 = nn.BatchNorm1d(200)\n",
    "        self.fc4 = nn.Linear(200, 2)\n",
    "        self.batchnorm4 = nn.BatchNorm1d(2)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.batchnorm1(F.relu(self.fc1(x)))\n",
    "        x = self.batchnorm2(F.relu(self.fc2(x)))\n",
    "        x = self.batchnorm3(F.relu(self.fc3(x)))\n",
    "        x = self.batchnorm4(self.fc4(x))\n",
    "        return F.log_softmax(x, dim=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Training methodology"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sklearn\n",
    "from sklearn.model_selection import KFold\n",
    "\n",
    "\n",
    "init_lr = 1e-1\n",
    "batch_size = 1024\n",
    "device = torch.device('cpu' if 'cpu' in memtorch.__version__ else 'cuda')\n",
    "\n",
    "def adjust_lr(optimizer, epoch):\n",
    "    lr = init_lr * (0.1 ** (epoch // 20))\n",
    "    for param_group in optimizer.param_groups:\n",
    "        param_group['lr'] = lr\n",
    "\n",
    "    return lr\n",
    "\n",
    "def train(net, train_loader, test_loader, epochs=10, fold=0):\n",
    "    print('fold %d' % fold)\n",
    "    best_f1_score = 0\n",
    "    for epoch in range(epochs):\n",
    "        lr = adjust_lr(optimizer, epoch)\n",
    "        running_loss = 0\n",
    "        for data in train_loader:\n",
    "            inputs, labels = data\n",
    "            inputs = inputs.float()\n",
    "            if device == torch.device('cuda'):\n",
    "                inputs = inputs.cuda()\n",
    "                labels = labels.cuda()\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "            outputs = net(inputs)\n",
    "            loss = criterion(outputs, labels)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            running_loss += loss.item()\n",
    "\n",
    "        f1_score = test(net, test_loader)\n",
    "        if f1_score > best_f1_score:\n",
    "            torch.save(net.state_dict(), 'trained_net_fold_%d.pt' % fold)\n",
    "            best_f1_score = f1_score\n",
    "\n",
    "        print('[Epoch %d] lr: %.4f f1_score: %.4f\\ttraining loss: %.4f' % (epoch + 1, lr, f1_score, running_loss / len(train_loader)))\n",
    "\n",
    "def test(net, test_loader):\n",
    "    confusion_matrix = torch.zeros(2, 2)\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    with torch.no_grad():\n",
    "        for data in test_loader:\n",
    "            inputs, labels = data\n",
    "            inputs = inputs.float()\n",
    "            if device == torch.device('cuda'):\n",
    "                inputs = inputs.cuda()\n",
    "                labels = labels.cuda()\n",
    "\n",
    "            outputs = net(inputs)\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            for t, p in zip(labels.view(-1), predicted.view(-1)):\n",
    "                confusion_matrix[t.long(), p.long()] += 1\n",
    "\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "\n",
    "    f1_score = 2 * confusion_matrix[0][0] / (2 * confusion_matrix[0][0] + confusion_matrix[0][1] + confusion_matrix[1][0])\n",
    "    return f1_score.item()\n",
    "\n",
    "dataset = SeizureDataset(csv_path)\n",
    "torch.manual_seed(0)\n",
    "kf = KFold(n_splits=5, shuffle=True)\n",
    "train_loaders = []\n",
    "test_loaders = []\n",
    "for i, (train_index, test_index) in enumerate(kf.split(dataset)):\n",
    "    train_ = torch.utils.data.Subset(dataset, train_index)\n",
    "    test_ = torch.utils.data.Subset(dataset, test_index)\n",
    "    train_loaders.append(torch.utils.data.DataLoader(train_, batch_size=batch_size, shuffle=True))\n",
    "    test_loaders.append(torch.utils.data.DataLoader(test_, batch_size=batch_size, shuffle=False))\n",
    "\n",
    "torch.manual_seed(torch.initial_seed())\n",
    "torch.save(test_loaders, 'test_loaders.pth')\n",
    "assert(len(train_loaders) == len(test_loaders))\n",
    "\n",
    "# Determine the baseline F1 score\n",
    "df = pd.DataFrame(columns=['indentifier', 'fold', 'f1_score'])\n",
    "for indentifier in range(100):\n",
    "    test_loaders = torch.load('test_loaders.pth')\n",
    "    for fold in range(len(test_loaders)):\n",
    "        net = EEGNet().to(device)\n",
    "        f1_score = test(net, test_loaders[fold])\n",
    "        df = df.append({'indentifier': indentifier, 'fold': fold, 'f1_score': f1_score}, ignore_index=True)\n",
    "\n",
    "    df.to_csv('baseline.csv', index=False)\n",
    "\n",
    "# Determine the F1 score\n",
    "for fold in range(len(train_loaders)):\n",
    "    net = EEGNet().to(device)\n",
    "    criterion = torch.nn.CrossEntropyLoss()\n",
    "    optimizer = torch.optim.Adam(net.parameters(), lr=init_lr)\n",
    "    train(net, train_loaders[fold], test_loaders[fold], epochs=50, fold=fold)\n",
    "\n",
    "fold_f1_scores = []\n",
    "for fold in range(len(train_loaders)):\n",
    "    net = EEGNet().to(device)\n",
    "    net.load_state_dict(torch.load('trained_net_fold_%d.pt' % fold), strict=True)\n",
    "    fold_f1_scores.append(test(net, test_loaders[fold]))\n",
    "    print('f1_score of fold %d: %0.4f' % (fold, fold_f1_scores[fold]))\n",
    "\n",
    "print('baseline -> mean: %0.4f\\tstddev: %0.4f' % (np.mean(df['f1_score'].values), np.std(df['f1_score'].values)))\n",
    "print('trained  -> mean: %0.4f\\tstddev: %0.4f' % (np.mean(fold_f1_scores), np.std(fold_f1_scores)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Network conversion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from memtorch.mn.Module import patch_model\n",
    "from memtorch.map.Parameter import naive_map\n",
    "from memtorch.bh.crossbar.Program import naive_program\n",
    "from memtorch.bh.nonideality.NonIdeality import apply_nonidealities\n",
    "import copy\n",
    "\n",
    "\n",
    "test_loaders = torch.load('test_loaders.pth')\n",
    "reference_memristor = memtorch.bh.memristor.VTEAM\n",
    "reference_memristor_params = {'time_series_resolution': 1e-6,\n",
    "                              'alpha_off': 1,\n",
    "                              'alpha_on': 3,\n",
    "                              'v_off': 0.5,\n",
    "                              'v_on': -0.53,\n",
    "                              'r_off': 2.5e3,\n",
    "                              'r_on': 100,\n",
    "                              'k_off': 4.03e-8,\n",
    "                              'k_on': -80,\n",
    "                              'd': 10e-9,\n",
    "                              'x_on': 0,\n",
    "                              'x_off': 10e-9}\n",
    "    \n",
    "# Determine the first baseline F1 score\n",
    "df = pd.DataFrame(columns=['indentifier', 'fold', 'f1_score'])\n",
    "for indentifier in range(100):\n",
    "    for fold in range(len(test_loaders)):\n",
    "        net = EEGNet()\n",
    "        net.load_state_dict(torch.load('trained_net_fold_%d.pt' % fold), strict=True).to(device)\n",
    "        patched_net = patch_model(copy.deepcopy(net),\n",
    "                                  memristor_model=reference_memristor,\n",
    "                                  memristor_model_params=reference_memristor_params,\n",
    "                                  module_parameters_to_patch=[torch.nn.Linear],\n",
    "                                  mapping_routine=naive_map,\n",
    "                                  transistor=True,\n",
    "                                  programming_routine=None,\n",
    "                                  scheme=memtorch.bh.Scheme.DoubleColumn)\n",
    "\n",
    "        for i, (name, m) in enumerate(list(patched_net.named_modules())):\n",
    "            if isinstance(m, memtorch.mn.Linear):\n",
    "                m.crossbars[0].conductance_matrix = m.crossbars[0].conductance_matrix.uniform_(1 / 2.5e3, 1 / 100)\n",
    "                m.crossbars[1].conductance_matrix = m.crossbars[1].conductance_matrix.uniform_(1 / 2.5e3, 1 / 100)\n",
    "\n",
    "        patched_net.tune_()\n",
    "        f1_score = test(patched_net, test_loaders[fold])\n",
    "        df = df.append({'indentifier': indentifier, 'fold': fold, 'f1_score': f1_score}, ignore_index=True)\n",
    "\n",
    "    df.to_csv('baseline_variability.csv', index=False)\n",
    "\n",
    "# Determine the second baseline F1 score\n",
    "df_2 = pd.DataFrame(columns=['indentifier', 'fold', 'f1_score'])\n",
    "for indentifier in range(100):\n",
    "    for fold in range(len(test_loaders)):\n",
    "        net = EEGNet().to(device)\n",
    "        patched_net = patch_model(copy.deepcopy(net),\n",
    "                                  memristor_model=reference_memristor,\n",
    "                                  memristor_model_params=reference_memristor_params,\n",
    "                                  module_parameters_to_patch=[torch.nn.Linear],\n",
    "                                  mapping_routine=naive_map,\n",
    "                                  transistor=True,\n",
    "                                  programming_routine=None,\n",
    "                                  scheme=memtorch.bh.Scheme.DoubleColumn)\n",
    "\n",
    "        for i, (name, m) in enumerate(list(patched_net.named_modules())):\n",
    "            if isinstance(m, memtorch.mn.Linear):\n",
    "                m.crossbars[0].conductance_matrix = m.crossbars[0].conductance_matrix.uniform_(1 / 2.5e3, 1 / 100)\n",
    "                m.crossbars[1].conductance_matrix = m.crossbars[1].conductance_matrix.uniform_(1 / 2.5e3, 1 / 100)\n",
    "\n",
    "        f1_score = test(patched_net, test_loaders[fold])\n",
    "        df_2 = df_2.append({'indentifier': indentifier, 'fold': fold, 'f1_score': f1_score}, ignore_index=True)\n",
    "\n",
    "    df_2.to_csv('baseline_variability_no_tune.csv', index=False)\n",
    "\n",
    "# Determine the F1 score\n",
    "fold_f1_scores = []\n",
    "for fold in range(len(test_loaders)):\n",
    "    net = EEGNet()\n",
    "    net.load_state_dict(torch.load('trained_net_fold_%d.pt' % fold), strict=False).to(device)\n",
    "    patched_net = patch_model(net,\n",
    "                              memristor_model=reference_memristor,\n",
    "                              memristor_model_params=reference_memristor_params,\n",
    "                              module_parameters_to_patch=[torch.nn.Linear],\n",
    "                              mapping_routine=naive_map,\n",
    "                              transistor=True,\n",
    "                              programming_routine=None,\n",
    "                              scheme=memtorch.bh.Scheme.DoubleColumn)\n",
    "\n",
    "    patched_net.tune_()\n",
    "    f1_score = test(patched_net, test_loaders[fold])\n",
    "    fold_f1_scores.append(f1_score)\n",
    "\n",
    "tuned_baseline = np.mean(df['f1_score'].values)\n",
    "print('baseline          -> mean: %0.4f\\tstddev: %0.4f' % (np.mean(df_2['f1_score'].values), np.std(df_2['f1_score'].values)))    \n",
    "print('tuned baseline    -> mean: %0.4f\\tstddev: %0.4f' % (np.mean(df['f1_score'].values), np.std(df['f1_score'].values)))\n",
    "print('trained and tuned -> mean: %0.4f\\tstddev: %0.4f' % (np.mean(fold_f1_scores), np.std(fold_f1_scores)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "# Plot the hysteresis loop\n",
    "hysteresis_loop_reference_memristor_params = {'time_series_resolution': 1e-6,\n",
    "                              'alpha_off': 1,\n",
    "                              'alpha_on': 3,\n",
    "                              'v_off': 0.5,\n",
    "                              'v_on': -0.53,\n",
    "                              'r_off': memtorch.bh.StochasticParameter(loc=2.5e3, scale=50, min=2),\n",
    "                              'r_on': memtorch.bh.StochasticParameter(loc=100, scale=25, min=1),\n",
    "                              'k_off': 4.03e-8,\n",
    "                              'k_on': -80,\n",
    "                              'd': 10e-9,\n",
    "                              'x_on': 0,\n",
    "                              'x_off': 10e-9}\n",
    "\n",
    "palette = [\"#DA4453\", \"#8CC152\", \"#4A89DC\", \"#F6BB42\", \"#B600B0\", \"#535353\"]\n",
    "f = plt.figure(figsize=(16/3, 4))\n",
    "plt.title('Hysteresis Loop')\n",
    "plt.xlabel('Voltage (V)')\n",
    "plt.ylabel('Current (A)')\n",
    "j = 0\n",
    "for i in range(10):\n",
    "    j = j + 1\n",
    "    if j == 6:\n",
    "        j = 0\n",
    "\n",
    "    memristor = reference_memristor(**hysteresis_loop_reference_memristor_params)\n",
    "    voltage_signal, current_signal = memristor.plot_hysteresis_loop(duration=2, voltage_signal_amplitude=1, voltage_signal_frequency = 0.5, return_result=True)\n",
    "    plt.plot(voltage_signal, current_signal, color=palette[j])\n",
    "\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Device-to-device variability investigation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Determine the F1 score\n",
    "df = pd.DataFrame(columns=['sigma', 'mean', 'std'])\n",
    "sigma_values = np.linspace(0, 500, 21)\n",
    "for sigma in sigma_values:\n",
    "    non_linear_reference_memristor_params = {'time_series_resolution': 1e-6,\n",
    "                              'alpha_off': 1,\n",
    "                              'alpha_on': 3,\n",
    "                              'v_off': 0.5,\n",
    "                              'v_on': -0.53,\n",
    "                              'r_off': memtorch.bh.StochasticParameter(loc=2.5e3, scale=sigma*2, min=1),\n",
    "                              'r_on': memtorch.bh.StochasticParameter(loc=100, scale=sigma, min=1),\n",
    "                              'k_off': 4.03e-8,\n",
    "                              'k_on': -80,\n",
    "                              'd': 10e-9,\n",
    "                              'x_on': 0,\n",
    "                              'x_off': 10e-9}\n",
    "    f1_scores = []\n",
    "    for fold in range(len(test_loaders)):\n",
    "        net = EEGNet()\n",
    "        net.load_state_dict(torch.load('trained_net_fold_%d.pt' % fold), strict=True).to(device)\n",
    "        patched_net = patch_model(copy.deepcopy(net),\n",
    "                                  memristor_model=reference_memristor,\n",
    "                                  memristor_model_params=non_linear_reference_memristor_params,\n",
    "                                  module_parameters_to_patch=[torch.nn.Linear],\n",
    "                                  mapping_routine=naive_map,\n",
    "                                  transistor=True,\n",
    "                                  programming_routine=None,\n",
    "                                  scheme=memtorch.bh.Scheme.DoubleColumn)\n",
    "\n",
    "        patched_net.tune_()\n",
    "        f1_score = test(patched_net, test_loaders[fold])\n",
    "        f1_scores.append(f1_score)\n",
    "\n",
    "    df = df.append({'sigma': sigma, 'mean': np.mean(f1_scores), 'std': np.std(f1_scores)}, ignore_index=True)\n",
    "\n",
    "df.to_csv('variability.csv', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = plt.figure(figsize=(16/3, 4))\n",
    "plt.axhline(y=tuned_baseline, color='k', linestyle='--', zorder=1)\n",
    "b = plt.bar(df['sigma'], df['mean'], width=12.5, zorder=2)\n",
    "plt.xlabel('$\\sigma$')\n",
    "plt.ylabel('F1 Score')\n",
    "for bar in b:\n",
    "    bar.set_edgecolor('black')\n",
    "    bar.set_facecolor(palette[0])\n",
    "    bar.set_linewidth(1)\n",
    "\n",
    "f.tight_layout()\n",
    "plt.grid()\n",
    "plt.ylim([0.9, 1.0])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Non-linear IV characteristics investigation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Determine the F1 score\n",
    "df = pd.DataFrame(columns=['sigma', 'mean', 'std'])\n",
    "sigma_values = np.linspace(0, 500, 11)\n",
    "f1_scores = []\n",
    "for fold in range(len(test_loaders)):\n",
    "    net = EEGNet()\n",
    "    net.load_state_dict(torch.load('trained_net_fold_%d.pt' % fold), strict=False).to(device)\n",
    "    patched_net = patch_model(net,\n",
    "                              memristor_model=reference_memristor,\n",
    "                              memristor_model_params=reference_memristor_params,\n",
    "                              module_parameters_to_patch=[torch.nn.Linear],\n",
    "                              mapping_routine=naive_map,\n",
    "                              transistor=True,\n",
    "                              programming_routine=None,\n",
    "                              scheme=memtorch.bh.Scheme.DoubleColumn)\n",
    "\n",
    "    patched_net = apply_nonidealities(patched_net,\n",
    "                    non_idealities=[memtorch.bh.nonideality.NonIdeality.NonLinear],\n",
    "                    sweep_duration=2,\n",
    "                    sweep_voltage_signal_amplitude=1,\n",
    "                    sweep_voltage_signal_frequency=0.5)\n",
    "\n",
    "    patched_net.tune_()\n",
    "    f1_score = test(patched_net, test_loaders[fold])\n",
    "    f1_scores.append(f1_score)\n",
    "\n",
    "print('mean: %0.4f\\tstddev: %0.4f' % (np.mean(f1_scores), np.std(f1_scores)))   "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
